import numpy as np
import pandas as pd
import nltk
from tqdm import tqdm

seed = 1024
np.random.seed(seed)

path = '../data/'

train = pd.read_pickle(path + "train_clean.pkl")
valid = pd.read_pickle(path + "valid_clean.pkl")
dev = pd.read_pickle(path+'dev_clean.pkl')


data_all = pd.concat([train,valid,dev])
data_all.reset_index(inplace=1,drop=1)



def get_pos_tag(cx):
    wl = str(cx).lower().split()
    pos_l = nltk.pos_tag(wl)
    q1_pos = []
    for pos in pos_l:
        q1_pos.append(pos[1])
    return q1_pos


# get_pos_tag(data_all['context'][1])
pos_fea = np.zeros((data_all.shape[0],2))
dd =data_all['context'].values
for it in tqdm(np.arange(data_all.shape[0])):
    pos_ = get_pos_tag(dd[it])
    pos_fea[it,0] = len(pos_)
    pos_fea[it,1] = len(set(pos_))


all_fea = pos_fea
train_fea = all_fea[:train.shape[0]]
valid_fea = all_fea[train.shape[0]:(train.shape[0]+valid.shape[0])]
dev_fea = all_fea[(train.shape[0]+valid.shape[0]):]

pd.to_pickle(train_fea,path+'train_pos.pkl')
pd.to_pickle(valid_fea,path+'valid_pos.pkl')
pd.to_pickle(dev_fea,path+'dev_pos.pkl')